#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <math.h>
#include "poisson.h"
#include "random.h"

/* log(n!) for n > 0 */
static double log_fact(int n)
{
  int i;
  double sum;
  
  sum = 0;
  for (i=2; i<=n; i++)
    sum += log((double) i);
  return sum;
}

static double poi_score(double mean, int val)
{
  if (val < 0)
    return -HUGE_VAL;
  
  return -mean + val * log(mean) - log_fact(val);
}

static double poi_prob(double mean, int val)
{
  return exp(poi_score(mean, val));
}

static double poi_sample(double mean)
{
  double desired;
  int n_plus;
  int n_minus;
  int delta;
  double prob_n_plus;
  double prob_n_minus;
  double cumprob;

  /* choose a random number */
  desired = RANDOM_DOUBLE;

  n_plus = n_minus = (int) mean;
  
  delta = 1;

  cumprob = prob_n_plus = prob_n_minus = poi_prob(mean, n_plus);

  while (cumprob < desired)
  {
    /* alternate between going above the mean and below until
     * it is not possible to go below */
    if (0 == n_minus)
      delta = 1;
    else
      delta *= -1;
    
    /* up move */
    if (delta > 0)
    {
      n_plus += 1;
      
      prob_n_plus *= mean / n_plus;
      
      cumprob += prob_n_plus;
    }
    else
    {
      n_minus -= 1;
      
      prob_n_minus *= ((double) n_minus+1) / mean;
      
      cumprob += prob_n_minus;
    }
  }
  
  /* at the end check if the last move was an up move or a down move */
  if (delta > 0)
    return n_plus;
  else
    return n_minus;

}

void Poisson_Init_Params(Poisson_t * dist, int nparams, ...)
{
  if (1 == nparams)
  {
    va_list ap;
  
    va_start(ap, nparams);

    dist->mean = va_arg(ap, double);
    dist->fixed_mean = 1;

    va_end(ap);

    if (dist->mean < 0)
    {
      fprintf(stderr, "error: Poisson requires a positive-valued parameter\n");
      exit(1);
    }
  }
  else if (0 == nparams)
  {
    dist->fixed_mean = 0;
  }
  else
  {
    fprintf(stderr, "error: Poisson takes at most 1 parameter\n");
    exit(1);
  }
}


int Poisson_Sample(const Poisson_t * dist, int nargs, ...)
{
  double mean;
  
  if (dist->fixed_mean)
    mean = dist->mean;

  else if (1 == nargs)
  {
    va_list ap;
  
    va_start(ap, nargs);

    mean = va_arg(ap, double);

    va_end(ap);
  }
  else
  {
    fprintf(stderr, "error: Poisson requires 1 argument\n");
    exit(1);
  }

  return poi_sample(mean);
}


double Poisson_LogProb(const Poisson_t * dist, int val, int nargs, ...)
{
  double mean;
  
  if (dist->fixed_mean)
    mean = dist->mean;

  else if (1 == nargs)
  {
    va_list ap;
  
    va_start(ap, nargs);

    mean = va_arg(ap, double);

    va_end(ap);
  }
  else
  {
    fprintf(stderr, "error: Poisson with 0 parameters requires 1 argument\n");
    exit(1);
  }

  return poi_score(mean, val);
}

double Poisson_Prob(const Poisson_t * dist, int val, int nargs, ...)
{
  double mean;
  
  if (dist->fixed_mean)
    mean = dist->mean;

  else if (1 == nargs)
  {
    va_list ap;
  
    va_start(ap, nargs);

    mean = va_arg(ap, double);

    va_end(ap);
  }
  else
  {
    fprintf(stderr, "error: Poisson with 0 parameters requires 1 argument\n");
    exit(1);
  }

  return exp(poi_score(mean, val));
  
}
